// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <cstdint>
#include <functional>
#include <memory>
#include <vector>

#include <turbo/utility/status.h>
#include <turbo/utility/status.h>

#include <aries/parquet/platform.h>

namespace nebula {
    class Array;
} // namespace nebula

namespace parquet {
    struct ArrowWriteContext;

    namespace nebula {
        // This files contain internal implementation details and should not be considered
        // part of the public API.

        // The MultipathLevelBuilder is intended to fully support all Arrow nested types that
        // map to parquet types (i.e. Everything but Unions).
        //

        /// \brief Half open range of elements in an array.
        struct ElementRange {
            /// Upper bound of range (inclusive)
            int64_t start;
            /// Upper bound of range (exclusive)
            int64_t end;

            bool Empty() const { return start == end; }

            int64_t Size() const { return end - start; }
        };

        /// \brief Result for a single leaf array when running the builder on the
        /// its root.
        struct MultipathLevelBuilderResult {
            /// \brief The Array containing only the values to write (after all nesting has
            /// been processed.
            ///
            /// No additional processing is done on this array (it is copied as is when
            /// visited via a DFS).
            std::shared_ptr<::nebula::Array> leaf_array;

            /// \brief Might be null.
            const int16_t *def_levels = nullptr;

            /// \brief  Might be null.
            const int16_t *rep_levels = nullptr;

            /// \brief Number of items (int16_t) contained in def/rep_levels when present.
            int64_t def_rep_level_count = 0;

            /// \brief Contains element ranges of the required visiting on the
            /// descendants of the final list ancestor for any leaf node.
            ///
            /// The algorithm will attempt to consolidate visited ranges into
            /// the smallest number possible.
            ///
            /// This data is necessary to pass along because after producing
            /// def-rep levels for each leaf array it is impossible to determine
            /// which values have to be sent to parquet when a null list value
            /// in a nullable ListArray is non-empty.
            ///
            /// This allows for the parquet writing to determine which values ultimately
            /// needs to be written.
            std::vector<ElementRange> post_list_visited_elements;

            /// Whether the leaf array is nullable.
            bool leaf_is_nullable;
        };

        /// \brief Logic for being able to write out nesting (rep/def level) data that is
        /// needed for writing to parquet.
        class PARQUET_EXPORT MultipathLevelBuilder {
        public:
            /// \brief A callback function that will receive results from the call to
            /// Write(...) below.  The MultipathLevelBuilderResult passed in will
            /// only remain valid for the function call (i.e. storing it and relying
            /// for its data to be consistent afterwards will result in undefined
            /// behavior.
            using CallbackFunction =
            std::function<::turbo::Status(const MultipathLevelBuilderResult &)>;

            /// \brief Determine rep/def level information for the array.
            ///
            /// The callback will be invoked for each leaf Array that is a
            /// descendant of array.  Each leaf array is processed in a depth
            /// first traversal-order.
            ///
            /// \param[in] array The array to process.
            /// \param[in] array_field_nullable Whether the algorithm should consider
            ///   the array column as nullable (as determined by its type's parent
            ///   field).
            /// \param[in, out] context for use when allocating memory, etc.
            /// \param[out] write_leaf_callback Callback to receive results.
            /// There will be one call to the write_leaf_callback for each leaf node.
            static ::turbo::Status Write(const ::nebula::Array &array, bool array_field_nullable,
                                         ArrowWriteContext *context,
                                         CallbackFunction write_leaf_callback);

            /// \brief Construct a new instance of the builder.
            ///
            /// \param[in] array The array to process.
            /// \param[in] array_field_nullable Whether the algorithm should consider
            ///   the array column as nullable (as determined by its type's parent
            ///   field).
            static ::turbo::Result<std::unique_ptr<MultipathLevelBuilder> > Make(
                const ::nebula::Array &array, bool array_field_nullable);

            virtual ~MultipathLevelBuilder() = default;

            /// \brief Returns the number of leaf columns that need to be written
            /// to Parquet.
            virtual int GetLeafCount() const = 0;

            /// \brief Calls write_leaf_callback with the MultipathLevelBuilderResult corresponding
            /// to |leaf_index|.
            ///
            /// \param[in] leaf_index The index of the leaf column to write.  Must be in the range
            /// [0, GetLeafCount()].
            /// \param[in, out] context for use when allocating memory, etc.
            /// \param[out] write_leaf_callback Callback to receive the result.
            virtual ::turbo::Status Write(int leaf_index, ArrowWriteContext *context,
                                          CallbackFunction write_leaf_callback) = 0;
        };
    } // namespace nebula
} // namespace parquet
